import mrjob
from mrjob.job import MRJob
from mrjob.step import MRStep
from mrjob import protocol
import json
import re


# class MyProtocol:
#     def read(self, bytes):
#         return "KEY",bytes.decode("utf-8")
#
#     def write(self, key, value):
#         return "{}\t{}".format(key, value).encode("utf-8")


class JobMRJob(MRJob):
    # INPUT_PROTOCOL = MyProtocol
    # INTERNAL_PROTOCOL = mrjob.protocol.JSONProtocol
    # OUTPUT_PROTOCOL = MyProtocol
    FILES = ["./lib.py"]

    def steps(self):
        return [
            MRStep(mapper=self.map_line,
                   reducer=self.reduce_jobs
                   ),
            MRStep(mapper=self.mapper2,
                   reducer=self.reducer2)
        ]

    def configure_args(self):
        super().configure_args()
        self.add_passthru_arg("-x",
                              type=int,
                              default=1,
                              choices=(1, 2),
                              help="specify the type to show")

    def map_line(self, key, value):  # None, line
        from .lib import get_distinct_and_avg_salary
        try:
            job_info = json.loads(value)
            yield get_distinct_and_avg_salary(job_info)
        except Exception as e:
            pass

    def reduce_jobs(self, key, values):
        values = list(values)

        avg_salary = sum(values) / len(values)
        yield (key, (avg_salary, len(values)))

    def mapper2(self, key, values):
        yield (None, (key, values))
        # yield None, values

    # (None, (('成都-青羊区', [14.5, 6]), ('成都-温江区', [14.5, 6])) ...)
    def reducer2(self, key, values):
        if self.options.x == 1:
            values = sorted(list(values), key=lambda e: e[1][1])[-3:]
        else:
            values = sorted(list(values), key=lambda e: e[1][0])[-3:]
        for value in values:
            yield value


if __name__ == "__main__":
    import os

    os.system("rm -rf ./result")
    JobMRJob.run()

    # for file in os.listdir("./result"):
    #     with open("./result/" + file, "r") as fp:
    #         while True:
    #             line = fp.readline()
    #             if not line:
    #                 break
    #
    #             key, value = line.split('\t')
    #             print("{}\t{}".format(json.loads(key), json.loads(value)))
